# -*- coding: utf-8 -*-
import cv2
import json
import os
import random
import logconfig
import glob
from xml.etree.ElementTree import Element, ElementTree

class PrepareDataset():
    '''
        准备数据集，主要工作有两个：
        1. 将YOLO json格式标签(labelme工具制作)转为VOC xml格式。其中json格式中的目标框格式由中心点坐标与宽高值组成, 
           而XML格式的目标框由左上角坐标与右下角坐标组成
        2. 把数据集拆分为训练集与测试集,比例由参数train_ration调整,默认0.8表示80%划分为训练集
    '''
    def __init__(self, train_ratio=0.8) -> None:
        self.logger = logconfig.init_log_config('prepare_dataset.log')
        self._load_config()  #读取项目根目录下的配置
        self.train_ratio = train_ratio
        self.missing = []    #数据集中无法找到的图片(可能打标签时失误造成)
        self.dataset = []    #图片与对应标签XML的集合
        self.label_stat = {} #标签数量统计

    def _load_config(self):
        '''
            加载配置文件,放到类全局的config变量
        '''
        with open('./config.json', 'r', encoding='utf8') as fp:
            self.config = json.load(fp)
            self.root = self.config['dataset']['root']
            self.path = os.path.join(self.root, self.config['dataset']['dir'])

    def draw_box(self, anno):
        '''
            把预测结果方框results画到原图中,便于观察效果
        '''
        image_fullpath = os.path.join(self.path, anno[0]['image'])
        if not os.path.exists(image_fullpath):
            return
        image = cv2.imread(image_fullpath)
        color = (0, 244, 10)
        thickness = 2
        font = cv2.FONT_HERSHEY_SIMPLEX
        for res in anno[0]['annotations']:
            label = res['label']
            box = res['coordinates']
            width = round(box['width'])
            height = round(box['height'])
            x = round(box['x'] - width/2)
            y = round(box['y'] - height/2)
            self.logger.debug('目标框左上坐标:{0}, 右下角坐标:{1}'.format((x,y), (x+width, y+height)))
            image_drawn = cv2.rectangle(
                image, (x, y), (x+width, y+height), color, thickness)
            # 标签位于检测框周围，避免超出窗口的情况（坐标为负）
            point = x, y - 10
            if point[1] < 20:
                point = point[0], y+height+22
            image_drawn = cv2.putText(image_drawn, '{0}'.format(label), point, font,
                                      1, color, thickness, cv2.LINE_8)
            self.show_window(anno[0]['image'], image_drawn)
        return image_drawn

    def show_window(self, name, image):
        '''
            使用图形窗口显示图片, 方便观察.如果没有接显示器请不要调用此方法.
        '''
        cv2.namedWindow(name, 1)
        cv2.moveWindow(name, 400, 200)
        cv2.imshow(name, image)
        cv2.waitKey(0)
        cv2.destroyWindow(name)
        cv2.waitKey(20)

    def view_some_images(self, jsons=[], limit=10):
        '''
        用图形窗口显示数据集中的数据，方便了解数据集。
        limit: 随机显示图片的数量（仅以jsons为空时有效）
        jsons: 指定要显示的图片标签文件名，例如['5051.json', '6001.json']
        '''
        if len(jsons) == 0:
            json_pattern = os.path.join(self.path, '*.json')
            jsons = random.choices(glob.glob(json_pattern), k=limit)
        else: #补充目录名
            jsons = [os.path.join(self.path, p) for p in jsons]
        for json_path in jsons:
            self.logger.debug('准备查看数据集:{0}'.format(json_path))
            with open(json_path, 'r') as fp:
                anno = json.load(fp)
                self.draw_box(anno)

    def prepare(self):        
        # 把所有标签写入label_list文件
        label_list=['background'] # 把0号标签当作背景

        json_pattern = os.path.join(self.path, '*.json')
        for json_path in glob.glob(json_pattern):
            label = self._json2xml(json_path)
            num = self.label_stat.get(label, 0)
            num += 1
            self.label_stat[label] = num
            if label is not None and label not in label_list:
                label_list.append(label)
        label_file = os.path.join(self.root, self.config['dataset']['labelfile'])
        with open(label_file, "w+", encoding='utf8') as fp:
            fp.write('\n'.join(label_list))
            self.logger.info('标签共{0}个类别，已写入文件:{1}'.format(len(label_list), label_file))
        self.logger.info('各类别图片统计情况：{0}'.format(self.label_stat))
        self.logger.info('成功转换图片标签共{0}个'.format(len(self.dataset)))

        if len(self.missing) > 0:
            with open(os.path.join(self.root, 'missing.txt'), 'w+', encoding='utf8') as fp:
                fp.write('\n'.join(self.missing))
            self.logger.warning('标签与图片不匹配共{0}个,详情保存于数据集的missing.txt'.format(len(self.missing)))
            
        self._split()

    def _split(self):
        random.shuffle(self.dataset)
        spliter = round(len(self.dataset) * self.train_ratio)
        trainset = self.dataset[0:spliter]
        evalset = self.dataset[spliter:]
        trainfile = os.path.join(self.root, self.config['dataset']['trainset'])
        evalfile  = os.path.join(self.root, self.config['dataset']['evalset'])
        with open(trainfile, 'w+', encoding='utf8') as fp:
            ts = '\n'.join([x[0] + ' ' + x[1] for x in trainset])
            fp.write(ts)
        with open(evalfile, 'w+', encoding='utf8') as fp:
            ts = '\n'.join([x[0] + ' ' + x[1] for x in evalset])
            fp.write(ts)

        self.logger.info('训练集共{0}张图，暂存到{1}'.format(len(trainset), trainfile))
        self.logger.info('验证集共{0}张图，暂存到{1}'.format(len(evalset), evalfile))


    def _json2xml(self, json_path):
        with open(json_path, 'r') as fp:
            anno = json.load(fp)
            image_path = os.path.join(self.config['dataset']['dir'], anno[0]['image'])
            image_fullpath = os.path.join(self.root, image_path)
            if not os.path.exists(image_fullpath):
                self.missing.append(anno[0]['image'])
                return
            #FIXME 假设一张图仅一个预测框，适用于当前比赛的简单任务
            res = anno[0]['annotations'][0]  
            label = res['label']  # 标签
            box = res['coordinates']
            width = round(box['width'])
            height = round(box['height'])
            # 将中心点坐标志为方框左上角坐标
            x = round(box['x'] - width/2)
            y = round(box['y'] - height/2)

            # 生成xml数据
            bbox = self._dict_to_xml('bndbox',
                                     {'xmin': x,
                                      'xmax': x+width, 
                                      'ymin': y, 
                                      'ymax': y+height})
            object_ = self._dict_to_xml(
                'object', {'name': label, 'difficult': 0}, [bbox])
            annotation = self._dict_to_xml('annotation', {"tmp": "tmp"}, [object_])
            # 根据xml数据生成xml文件格式
            tree = ElementTree(annotation)
            
            xml_file = os.path.splitext(anno[0]['image'])[0]+'.xml'
            # 为了兼容官方给的ssd_trainer.py，把Label写死
            label_dir = 'Label'
            label_fullpath = os.path.join(self.root, label_dir)
            if not os.path.exists(label_fullpath):
                os.mkdir(label_fullpath)
            xml_path = os.path.join(label_fullpath, xml_file)
            # 保存对应xml文件
            tree.write(xml_path, encoding='utf-8')
            #将图片文件名与xml保存到全局以便后面分割数据集
            self.dataset.append((image_path, os.path.join(label_dir, xml_file)))

            return label

    # 字典转xml元素数据
    def _dict_to_xml(self, tag, d, sub_elem=None):
        elem = Element(tag)
        for key, value in d.items():
            child = Element(key)
            child.text = str(value)
            elem.append(child)
        if sub_elem:
            for e in sub_elem:
                elem.append(e)
        return elem


if __name__ == "__main__":
    pd = PrepareDataset()
    pd.prepare()
    #随机看5张数据集中的图片
    # pd.view_some_images(limit=5)
    #指定查看数据集中的某张图片
    # pd.view_some_images(jsons=['5051.json'])
